package top.yudoge.dogeeditor.llm.rewrite

import cn.hutool.json.JSONUtil
import dev.langchain4j.model.chat.ChatLanguageModel
import dev.langchain4j.model.chat.StreamingChatLanguageModel
import top.yudoge.dogeeditor.llm.objects.query.RewriteQuery
import top.yudoge.dogeeditor.llm.util.PromptUtils
import java.io.BufferedReader
import java.io.InputStreamReader

class TaskChain(
    val nodes: List<TaskChainNode>,
) {

    val chainStore: MutableMap<String, Any> = mutableMapOf() // shared data storage for all nodes in a chain's execution

    var currentNodeIndex: Int = 0

    fun process(query: RewriteQuery, callback: RewriteCallback, model: StreamingChatLanguageModel, chatModel: ChatLanguageModel) {
        val ctx = ChainContext(
            query,
            callback,
            model,
            chatModel,
            this
        )
        currentNodeIndex = 0
        nodes.first().process(ctx)
    }
}

interface TaskChainNode {
    fun process(ctx: ChainContext)
    fun sendReasoningCallback(ctx: ChainContext, reasoningForWhat: String) {
        ctx.callback.reasoning(reasoningForWhat)
    }

    fun sendText(ctx: ChainContext, text: String) {
        ctx.callback.token("$text\n")
    }

    fun getChainStore(ctx: ChainContext, key: String): Any? {
        val ctxStore = ctx.taskChain.chainStore
        return ctxStore[key]
    }

    fun putChainStore(ctx: ChainContext, key: String, value: Any) {
        val ctxStore = ctx.taskChain.chainStore
        ctxStore[key] = value
    }

    fun removeChainStore(ctx: ChainContext, key: String) {
        val ctxStore = ctx.taskChain.chainStore
        ctxStore.remove(key)
    }

    fun moveBack(ctx: ChainContext) {
        ctx.taskChain.currentNodeIndex--
    }

    fun emitNext(ctx: ChainContext) {
        ctx.taskChain.currentNodeIndex++
        ctx.taskChain.nodes[ctx.taskChain.currentNodeIndex].process(ctx)
    }

    fun loadPrompt(resourcePath: String, model: Map<String, String>): String {
        return PromptUtils.loadPrompt(resourcePath, model)
    }

    fun safeReadJson(contentFromModel: String): String {
        return contentFromModel.replace("```json", "").replace("```", "")
    }

    fun <T> safeReadJsonArray(contentFromModel: String, clz: Class<T>): List<T> {
        var jsonArray = safeReadJson(contentFromModel)
        var parseArray = JSONUtil.parseArray(jsonArray).toList(clz)
        return parseArray
    }
}

class ChainContext(
    val query: RewriteQuery,
    val callback: RewriteCallback,
    val streamingModel: StreamingChatLanguageModel,
    val chatModel: ChatLanguageModel,
    val taskChain: TaskChain
)
    